package leetcode_0992

/**
 * TODO:改进方法暂时还是有问题
 */

// 会超时
func subarraysWithKDistinct(A []int, K int) int {

	res := 0

	length := len(A)
	if length < K {
		return 0
	}

	// 使用滑动窗口（逐个递增）
	for i := K; i <= length; i++ {
		res += slideWindow(A, i, K)
	}

	return res
}

func slideWindow(arr []int, windowLen int, k int) int {

	count := 0

	tmpMap := map[int]int{}
	for i := 0; i < windowLen; i++ {
		tmpMap[arr[i]] = 1
	}

	if len(tmpMap) == k {
		count++
	}

	for i := 1; i <= len(arr)-windowLen; i++ {

		// 删除元素
		if value, ok := tmpMap[arr[i-1]]; ok && value == 1 {
			delete(tmpMap, arr[i-1])
		} else {
			tmpMap[arr[i-1]]--
		}

		// 添加元素
		if _, ok := tmpMap[arr[i+windowLen-1]]; ok {
			tmpMap[arr[i+windowLen-1]]++
		} else {
			tmpMap[arr[i+windowLen-1]] = 1
		}

		// 是否满足条件
		if len(tmpMap) == k {
			count++
		}
	}

	return count
}

// 改进方法
func subarraysWithKDistinct1(A []int, K int) int {
	res := 0

	length := len(A)

	// 初始化
	tmpMap := map[int]int{}
	for j := 0; j < K; j++ {
		if _, ok := tmpMap[A[j]]; ok {
			tmpMap[A[j]]++
		} else {
			tmpMap[A[j]] = 1
		}
	}

	for i := 0; i <= len(A)-K; i++ {

		lenMap := len(tmpMap)
		if lenMap == K {
			res++
		}

		for j := i + K; j < length; j++ {
			if _, ok := tmpMap[A[j]]; !ok {
				tmpMap[A[j]] = 1
				lenMap++
			}

			if lenMap == K {
				res++
			} else if lenMap > K {
				break
			}
		}

	}

	return res
}
